{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9d9063",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Input, Dense, BatchNormalization\n",
    "from keras.regularizers import L1L2\n",
    "from keras.optimizers import Adam\n",
    "\n",
    "# customized metrics\n",
    "# out-of-sample r squared for keras\n",
    "def R_oos_tf(y_true, y_pred):\n",
    "    resid = tf.square(y_true-y_pred)\n",
    "    denom = tf.square(y_true)\n",
    "    return 1 - tf.divide(tf.reduce_mean(resid),tf.reduce_mean(denom))\n",
    "\n",
    "# data standardization\n",
    "# please standardize the data if BatchNormalization is not used\n",
    "def standardize(X_trn, X_vld, X_tst):\n",
    "    mu_trn = np.mean(np.array(X_trn),axis=0).reshape((1,X_trn.shape[1]))\n",
    "    sigma_trn = np.std(np.array(X_trn),axis=0).reshape((1,X_trn.shape[1]))\n",
    "\n",
    "    X_trn_std = (np.array(X_trn)-mu_trn)/sigma_trn\n",
    "    X_vld_std = (np.array(X_vld)-mu_trn)/sigma_trn\n",
    "    X_tst_std = (np.array(X_tst)-mu_trn)/sigma_trn\n",
    "    return X_trn_std, X_vld_std, X_tst_std\n",
    "\n",
    "# NN class\n",
    "class NN:\n",
    "    \n",
    "    def __init__(\n",
    "        self, n_layers=1, loss='mse', l1=1e-5, l2=0, learning_rate=.01, BatchNormalization=True, patience=5,\n",
    "        epochs=100, batch_size=3000, verbose=1, random_state=12308, monitor='val_R_oos_tf', base_neurons=5\n",
    "    ):\n",
    "        self.n_layers = n_layers\n",
    "        self.l1 = l1\n",
    "        self.l2 = l2\n",
    "        self.learning_rate = learning_rate\n",
    "        self.BatchNormalization = BatchNormalization\n",
    "        self.patience = patience\n",
    "        self.epochs = epochs\n",
    "        self.batch_size = batch_size\n",
    "        self.verbose = verbose\n",
    "        self.random_state = random_state\n",
    "        self.monitor = monitor\n",
    "        self.base_neurons = base_neurons\n",
    "\n",
    "    def set_params(self, **params):\n",
    "        for param in params.keys():\n",
    "            setattr(self, param, params[param])\n",
    "        return self\n",
    "    \n",
    "    def fit(self, X_trn, y_trn, X_vld, y_vld):\n",
    "        # fix random seed for reproducibility\n",
    "        random.seed(self.random_state)\n",
    "        np.random.seed(self.random_state)\n",
    "        tf.random.set_seed(self.random_state)\n",
    "        \n",
    "        # model construction\n",
    "        mod = Sequential()\n",
    "        mod.add(Input(shape=(X_trn.shape[1],)))\n",
    "        \n",
    "        for i in np.arange(self.n_layers,0,-1):\n",
    "            if self.n_layers>self.base_neurons:\n",
    "                if self.n_layers == i:\n",
    "                    mod.add(Dense(2**i, activation='relu'))\n",
    "                else:\n",
    "                    mod.add(Dense(2**i, activation='relu', kernel_regularizer=L1L2(self.l1,self.l2)))\n",
    "            else:\n",
    "                if self.n_layers == i:\n",
    "                    mod.add(Dense(2**(self.base_neurons-(self.n_layers-i)), activation='relu'))\n",
    "                else:\n",
    "                    mod.add(Dense(2**(self.base_neurons-(self.n_layers-i)), \n",
    "                                  activation='relu', kernel_regularizer=L1L2(self.l1,self.l2)))\n",
    "            if self.BatchNormalization:\n",
    "                mod.add(BatchNormalization())\n",
    "        \n",
    "        mod.add(Dense(1, kernel_regularizer=L1L2(self.l1,self.l2)))\n",
    "        \n",
    "        # early stopping\n",
    "        earlystop = tf.keras.callbacks.EarlyStopping(monitor=self.monitor, patience=self.patience,mode='max')\n",
    "\n",
    "        # Adam solver\n",
    "        opt = Adam(learning_rate=self.learning_rate)\n",
    "        \n",
    "        # compile the model\n",
    "        mod.compile(loss=self.loss,\n",
    "                    optimizer=opt,\n",
    "                    metrics=[R_oos_tf])\n",
    "\n",
    "        # fit the model\n",
    "        mod.fit(X_trn, np.array(y_trn).reshape((len(y_trn),1)), epochs=self.epochs, batch_size=self.batch_size, \n",
    "                callbacks=[earlystop], verbose=self.verbose, \n",
    "                validation_data=(X_vld,np.array(y_vld).reshape((len(y_vld),1))))\n",
    "        \n",
    "        self.model = mod\n",
    "        return self\n",
    "    \n",
    "    def predict(self, X):\n",
    "        return self.model.predict(X, verbose=self.verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2430701",
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {\n",
    "    'n_layers': [1],\n",
    "    'loss': ['mse'],\n",
    "    'l1': [1e-5,1e-3],\n",
    "    'learning_rate': [.001,.01],\n",
    "    'batch_size': [int(X_trn.shape[0]/50)],\n",
    "    'epochs': [100],\n",
    "    'random_state': [12308],\n",
    "    'BatchNormalization': [True],\n",
    "    'patience':[5],\n",
    "    'verbose': [0],\n",
    "    'monitor':['val_loss','val_R_oos_tf']\n",
    "}\n",
    "NN1 = val_fun(NN,params=params,X_trn=X_trn,y_trn=y_trn,X_vld=X_vld,y_vld=y_vld,is_NN=True,sleep=5)\n",
    "evaluate(y_trn, NN1.predict(X_trn), insample=True) \n",
    "evaluate(y_tst, NN1.predict(X_tst))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fhuan",
   "language": "python",
   "name": "fhuan"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
